// Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "memento_common.h"
#include "recovery_installer.h"

#include <iostream>
#include <fstream>
#include <string>
#include <sstream>

#include <unistd.h>

#include "error.h"

using std::cerr;
using std::endl;
using std::string;
using std::stringstream;

namespace chromeos_memento_updater {


const string RecoveryInstaller::kConfigPath = "recovery.conf";
const string RecoveryInstaller::kZipImagePath = "release_image.bin.zip";
const string RecoveryInstaller::kConfigUrl =
    "https://dl.google.com/dl/edgedl/chromeos/recovery/recovery.conf";


RecoveryInstaller::RecoveryInstaller(const string& running_dir,
                                     const string& board,
                                     const string& image_dir)
    : DiskImageInstaller(running_dir, board, "", kZip) {
  image_dir_ = image_dir + "/";
  image_path_ = image_dir_ + kZipImagePath;
  config_path_ = image_dir_ + kConfigPath;
  DownloadImageFromServer();
}

// If image_dir is not provided, we use /tmp as the default tmp directory.
RecoveryInstaller::RecoveryInstaller(const string& running_dir,
                                     const string& board)
    : DiskImageInstaller(running_dir, board, "", kZip) {
  image_dir_ = "/tmp/";
  image_path_ = image_dir_ + kZipImagePath;
  config_path_ = image_dir_ + kConfigPath;
  DownloadImageFromServer();
}

RecoveryInstaller::~RecoveryInstaller() {
  remove(image_path_.c_str());
  remove(config_path_.c_str());
}

void RecoveryInstaller::InstallDevice(const int channel_id,
                                      const string& device_path,
                                      const string& kernel_device_path) {
  //TODO(chunyen): Replace these with memento_common.h
  const int kFactoryKernelPartition = 2;
  const int kFactoryPartition = 3;
  const int kReleasePartition = 5;

  // For recovery image, real kernel is at partition 4 and rootfs is at
  // partition 3. So we need to change kReleasePartition to channel 3 and do
  // PostprocessKernel() for kReleaseKernelPartition.
  // And there is no Factory partitions in recovery image.

  switch(channel_id) {
    case kFactoryKernelPartition:
    case kFactoryPartition:
      error_exit("There is no factory image in recovery mode");
      break;
    case kReleasePartition:
      DiskImageInstaller::InstallDevice(3, device_path,
                                       kernel_device_path);
    default:
      DiskImageInstaller::InstallDevice(channel_id, device_path,
                                       kernel_device_path);
  }
}


bool RecoveryInstaller::FindKeyValue(const string& line, const string& key,
                                     string& value) {
  if (line.find(key) == 0) {
    int equal_pos = key.length();
    if (line[equal_pos] == '=') {
      value = line.substr(equal_pos + 1);
      return true;
    }
  }
  return false;
}

void RecoveryInstaller::DownloadImageFromServer() {
  int retry = 3;
  // To avoid unstable internet, retry multiple times here.
  while (retry-- > 0) {
    if (system(("wget " + kConfigUrl +
                " -O " + config_path_ +
                " --no-check-certificate").c_str()) == 0) {
      // Succeed and break;
      break;
    }
    cerr << "Connection failed, wait 3 seconds and retry(" << retry <<
            "times left)." << endl;
    sleep(3);
  }

  if (retry == 0) {
    error_exit("Fail downloading recovery image configuration.");
  }

  // To avoid the confusion of x86-alex and x86-alex-he,
  // change from "board" to "board_"
  std::ifstream fin;
  string buffer, filename, download_link;
  string board_underscore = board_ + "_";

  fin.open((config_path_).c_str());
  while(std::getline(fin, buffer)) {
    if (FindKeyValue(buffer, "url", filename)) {
      if (filename.find(board_) != string::npos) {
        download_link = filename;
        break;
      }
    }
  }

  // Can not find download url, error
  if (download_link.empty()) {
    error_exit("Can not find download url, installation failed.");
  }

  // Download the release image, retry 3 times
  retry = 3;
  while (retry-- > 0) {
    if (system(("wget " + download_link +
                " -O " + image_path_ +
                " --no-check-certificate").c_str()) == 0) {
      // Succeed, return
      return;
    }
  }
  // Failed to download image
  error_exit("Fail downloading release image.");
}

}  // namespace chromeos_memento_updater
